package com.mm.cloth.rmbg;

import ai.onnxruntime.*;
import android.content.res.AssetManager;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.util.Log;

import java.io.InputStream;
import java.nio.FloatBuffer;
import java.util.Collections;

/**
 * 基于ONNX Runtime的LaMa推理助手
 * 用于图像修复和去水印功能
 */
public class LamaInferenceHelper {
    private static final String TAG = "LamaInferenceHelper";
    
    private OrtSession session;
    private OrtEnvironment env;
    
    // LaMa模型参数
    private static final String INPUT_NAME = "input";
    private static final String OUTPUT_NAME = "output";
    
    // 图像标准化参数（LaMa通常使用简单的[0,1]标准化）
    private static final float[] MEAN = {0.5f, 0.5f, 0.5f};
    private static final float[] STD = {0.5f, 0.5f, 0.5f};
    
    /**
     * 初始化ONNX Runtime环境
     * @param assets AssetManager
     * @param modelPath 模型文件路径
     */
    public void initModel(AssetManager assets, String modelPath) throws Exception {
        Log.d(TAG, "初始化Big-LaMa FP16 ONNX Runtime环境，模型路径: " + modelPath);

        try {
            // 检查ONNX Runtime是否可用
            try {
                Class.forName("ai.onnxruntime.OrtEnvironment");
                Log.d(TAG, "ONNX Runtime类可用");
            } catch (ClassNotFoundException e) {
                throw new Exception("ONNX Runtime不可用", e);
            }

            // 检查系统架构兼容性
            String arch = System.getProperty("os.arch");
            Log.d(TAG, "系统架构: " + arch);

            // 创建ONNX Runtime环境（使用更安全的方式）
            try {
                env = OrtEnvironment.getEnvironment();
                Log.d(TAG, "ONNX Runtime环境创建成功");
            } catch (Exception e) {
                Log.e(TAG, "ONNX Runtime环境创建失败", e);
                throw new Exception("ONNX Runtime环境初始化失败: " + e.getMessage(), e);
            }

            // 配置会话选项（针对200MB大模型优化）
            OrtSession.SessionOptions options = new OrtSession.SessionOptions();
            try {
                // 检查可用内存，动态调整配置
                Runtime runtime = Runtime.getRuntime();
                long availableMemory = runtime.maxMemory() - (runtime.totalMemory() - runtime.freeMemory());
                long availableMB = availableMemory / (1024 * 1024);

                Log.d(TAG, "配置会话选项，可用内存: " + availableMB + "MB");

                if (availableMB < 200) {
                    // 内存不足，使用保守配置
                    Log.w(TAG, "内存不足，使用保守配置");
                    options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.NO_OPT);
                    options.setIntraOpNumThreads(1);
                    options.setInterOpNumThreads(1);
                } else if (availableMB < 400) {
                    // 内存一般，使用基础配置
                    Log.d(TAG, "内存一般，使用基础配置");
                    options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.BASIC_OPT);
                    options.setIntraOpNumThreads(2);
                    options.setInterOpNumThreads(1);
                } else {
                    // 内存充足，使用优化配置
                    Log.d(TAG, "内存充足，使用优化配置");
                    options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.BASIC_OPT);
                    int cpuCores = Runtime.getRuntime().availableProcessors();
                    int threads = Math.min(cpuCores, 4);
                    options.setIntraOpNumThreads(threads);
                    options.setInterOpNumThreads(2);
                }

                // 设置内存模式为保守模式
                try {
                    options.addConfigEntry("session.memory_pattern", "0"); // 保守内存模式
                    options.addConfigEntry("session.use_env_allocators", "1"); // 使用环境分配器
                } catch (Exception e) {
                    Log.w(TAG, "设置内存配置失败，继续使用默认配置", e);
                }

                Log.d(TAG, "Big-LaMa会话选项配置完成");
            } catch (Exception e) {
                Log.e(TAG, "会话选项配置失败", e);
                throw new Exception("会话选项配置失败: " + e.getMessage(), e);
            }

            // 检查模型文件
            InputStream modelStream = assets.open(modelPath);
            int available = modelStream.available();
            float modelSizeMB = available / (1024.0f * 1024.0f);
            Log.d(TAG, "Big-LaMa模型文件大小: " + String.format("%.1f", modelSizeMB) + " MB");

            if (available == 0) {
                modelStream.close();
                throw new Exception("模型文件为空");
            }

            // 验证是否为预期的大模型文件
            if (available < 150 * 1024 * 1024) { // 小于150MB
                Log.w(TAG, "⚠️ 模型文件小于150MB，可能不是big-lama_fp16.onnx");
            } else if (available > 250 * 1024 * 1024) { // 大于250MB
                Log.w(TAG, "⚠️ 模型文件大于250MB，可能内存不足");
            } else {
                Log.d(TAG, "✅ 模型文件大小符合预期的Big-LaMa FP16模型");
            }

            // 分块加载模型数据以避免内存峰值
            Log.d(TAG, "开始分块加载大模型数据...");
            byte[] modelBytes = new byte[available];
            int totalRead = 0;
            int chunkSize = 1024 * 1024; // 1MB chunks

            while (totalRead < available) {
                int remainingBytes = available - totalRead;
                int readSize = Math.min(chunkSize, remainingBytes);
                int read = modelStream.read(modelBytes, totalRead, readSize);

                if (read == -1) break;
                totalRead += read;

                // 显示加载进度
                if (totalRead % (10 * 1024 * 1024) == 0 || totalRead == available) {
                    float progress = (totalRead * 100.0f) / available;
                    Log.d(TAG, "模型加载进度: " + String.format("%.1f", progress) + "%");
                }
            }
            modelStream.close();

            Log.d(TAG, "Big-LaMa模型数据加载完成: " + (totalRead / (1024*1024)) + " MB");

            // 验证模型数据完整性
            if (totalRead != available) {
                throw new Exception("模型文件读取不完整: " + totalRead + "/" + available + " bytes");
            }

            if (totalRead < 100 * 1024 * 1024) { // 小于100MB肯定有问题
                throw new Exception("模型文件太小，可能损坏或不是正确的big-lama_fp16.onnx文件");
            }

            // 创建会话（针对200MB大模型优化）
            try {
                Log.d(TAG, "开始创建ONNX会话...");
                Log.d(TAG, "模型大小: " + (modelBytes.length / (1024*1024)) + "MB，预计需要5-15秒");

                // 再次强制垃圾回收，确保有足够内存
                System.gc();

                // 使用异步方式创建会话，避免ANR
                long sessionStartTime = System.currentTimeMillis();

                // 创建会话
                session = env.createSession(modelBytes, options);

                long sessionEndTime = System.currentTimeMillis();
                Log.d(TAG, "ONNX会话创建成功，耗时: " + (sessionEndTime - sessionStartTime) + "ms");

                // 验证会话是否可用
                if (session == null) {
                    throw new Exception("会话创建返回null");
                }

                // 验证会话功能
                try {
                    Log.d(TAG, "验证会话功能...");
                    int inputCount = session.getInputNames().size();
                    int outputCount = session.getOutputNames().size();
                    Log.d(TAG, "会话验证成功 - 输入: " + inputCount + ", 输出: " + outputCount);
                } catch (Exception e) {
                    Log.e(TAG, "会话功能验证失败", e);
                    throw new Exception("会话功能验证失败: " + e.getMessage(), e);
                }

            } catch (OutOfMemoryError e) {
                Log.e(TAG, "创建ONNX会话时内存不足", e);
                // 强制垃圾回收
                System.gc();
                throw new Exception("内存不足，无法加载200MB模型。请关闭其他应用后重试", e);
            } catch (Exception e) {
                Log.e(TAG, "ONNX会话创建失败", e);
                String errorMsg = "ONNX会话创建失败";
                if (e.getMessage() != null) {
                    if (e.getMessage().contains("OutOfMemoryError") || e.getMessage().contains("memory")) {
                        errorMsg = "内存不足，无法加载200MB模型";
                    } else if (e.getMessage().contains("timeout")) {
                        errorMsg = "模型加载超时，请重试";
                    } else {
                        errorMsg = "会话创建失败: " + e.getMessage();
                    }
                }
                throw new Exception(errorMsg, e);
            }

            // 打印模型信息
            printModelInfo();

        } catch (Exception e) {
            Log.e(TAG, "LaMa ONNX Runtime模型初始化失败", e);
            // 清理资源
            if (session != null) {
                try { session.close(); } catch (Exception ignored) {}
                session = null;
            }
            if (env != null) {
                try { env.close(); } catch (Exception ignored) {}
                env = null;
            }
            throw e;
        }
    }
    
    /**
     * 打印模型信息
     */
    private void printModelInfo() {
        try {
            Log.d(TAG, "=== LaMa模型信息 ===");
            Log.d(TAG, "输入节点数量: " + session.getInputNames().size());
            Log.d(TAG, "输出节点数量: " + session.getOutputNames().size());
            
            for (String inputName : session.getInputNames()) {
                Log.d(TAG, "输入节点: " + inputName);
            }
            
            for (String outputName : session.getOutputNames()) {
                Log.d(TAG, "输出节点: " + outputName);
            }
            
        } catch (Exception e) {
            Log.w(TAG, "获取模型信息失败", e);
        }
    }
    
    /**
     * 执行图像修复推理
     * @param originalImage 原始图像
     * @param maskImage 遮罩图像（白色区域为需要修复的区域）
     * @return 修复后的图像
     */
    public Bitmap inpaint(Bitmap originalImage, Bitmap maskImage) throws Exception {
        if (session == null || env == null) {
            throw new IllegalStateException("模型未初始化");
        }

        if (originalImage == null || maskImage == null) {
            throw new IllegalArgumentException("输入图像或遮罩不能为空");
        }

        Log.d(TAG, "开始LaMa图像修复推理");
        Log.d(TAG, "原始图像尺寸: " + originalImage.getWidth() + "x" + originalImage.getHeight());
        Log.d(TAG, "遮罩图像尺寸: " + maskImage.getWidth() + "x" + maskImage.getHeight());

        long startTime = System.currentTimeMillis();
        OnnxTensor tensor = null;

        try {
            // 1. 预处理：将图像和遮罩合并为4通道输入
            float[] inputData = preprocessImageWithMask(originalImage, maskImage);
            long[] shape = {1, 4, originalImage.getHeight(), originalImage.getWidth()}; // NCHW格式

            Log.d(TAG, "预处理完成，输入数据长度: " + inputData.length + ", 形状: [" +
                  shape[0] + "," + shape[1] + "," + shape[2] + "," + shape[3] + "]");

            // 2. 构建输入张量
            try {
                Log.d(TAG, "创建输入张量...");
                tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(inputData), shape);
                Log.d(TAG, "输入张量创建成功");
            } catch (Exception e) {
                Log.e(TAG, "输入张量创建失败", e);
                throw new Exception("输入张量创建失败: " + e.getMessage(), e);
            }

            // 3. 执行推理
            long inferenceStart = System.currentTimeMillis();
            Log.d(TAG, "开始ONNX推理...");
            try (OrtSession.Result result = session.run(Collections.singletonMap(INPUT_NAME, tensor))) {
                long inferenceEnd = System.currentTimeMillis();
                Log.d(TAG, "ONNX推理完成，耗时: " + (inferenceEnd - inferenceStart) + "ms");

                // 4. 获取输出结果
                try {
                    Log.d(TAG, "获取推理结果...");
                    Object outputObject = result.get(0).getValue();
                    if (outputObject == null) {
                        throw new Exception("推理结果为null");
                    }

                    Log.d(TAG, "转换输出数据...");
                    float[] outputData = convertOutputToFloatArray(outputObject);
                    if (outputData == null || outputData.length == 0) {
                        throw new Exception("输出数据为空");
                    }

                    // 5. 后处理：将输出转换为Bitmap
                    Log.d(TAG, "开始后处理...");
                    Bitmap resultBitmap = postprocessOutput(outputData, originalImage.getWidth(), originalImage.getHeight());
                    if (resultBitmap == null) {
                        throw new Exception("后处理失败，结果为null");
                    }

                    Log.d(TAG, "后处理完成");
                    return resultBitmap;

                } catch (Exception e) {
                    Log.e(TAG, "输出处理失败", e);
                    throw new Exception("输出处理失败: " + e.getMessage(), e);
                }


            }

        } catch (OutOfMemoryError e) {
            Log.e(TAG, "LaMa推理内存不足", e);
            System.gc();
            throw new Exception("内存不足，请尝试使用较小的图片", e);
        } catch (Exception e) {
            Log.e(TAG, "LaMa推理失败", e);
            throw e;
        } finally {
            // 确保释放张量资源
            if (tensor != null) {
                try {
                    tensor.close();
                } catch (Exception e) {
                    Log.w(TAG, "释放张量资源失败", e);
                }
            }
        }
    }
    
    /**
     * 预处理图像和遮罩
     * @param originalImage 原始图像
     * @param maskImage 遮罩图像
     * @return 4通道输入数据 (RGB + Mask)
     */
    private float[] preprocessImageWithMask(Bitmap originalImage, Bitmap maskImage) {
        int width = originalImage.getWidth();
        int height = originalImage.getHeight();
        
        // 确保遮罩尺寸与原图一致
        Bitmap resizedMask = maskImage;
        if (maskImage.getWidth() != width || maskImage.getHeight() != height) {
            resizedMask = Bitmap.createScaledBitmap(maskImage, width, height, true);
        }
        
        // 获取像素数据
        int[] originalPixels = new int[width * height];
        int[] maskPixels = new int[width * height];
        
        originalImage.getPixels(originalPixels, 0, width, 0, 0, width, height);
        resizedMask.getPixels(maskPixels, 0, width, 0, 0, width, height);
        
        // 转换为NCHW格式的float数组 (4通道: RGB + Mask)
        int totalPixels = width * height;
        float[] inputData = new float[4 * totalPixels];
        
        // 分别处理R、G、B、Mask四个通道
        for (int c = 0; c < 4; c++) {
            for (int h = 0; h < height; h++) {
                for (int w = 0; w < width; w++) {
                    int pixelIndex = h * width + w;
                    
                    float value;
                    if (c < 3) { // RGB通道
                        int pixel = originalPixels[pixelIndex];
                        if (c == 0) { // R通道
                            value = Color.red(pixel) / 255.0f;
                        } else if (c == 1) { // G通道
                            value = Color.green(pixel) / 255.0f;
                        } else { // B通道
                            value = Color.blue(pixel) / 255.0f;
                        }
                        
                        // 标准化到[-1,1]范围: (value - 0.5) / 0.5 = value * 2 - 1
                        value = value * 2.0f - 1.0f;
                    } else { // Mask通道
                        int maskPixel = maskPixels[pixelIndex];
                        // 将遮罩转换为0-1值（白色=1，黑色=0）
                        value = Color.red(maskPixel) / 255.0f;
                    }
                    
                    // NCHW格式存储
                    int outputIndex = c * totalPixels + pixelIndex;
                    inputData[outputIndex] = value;
                }
            }
        }
        
        // 清理临时bitmap
        if (resizedMask != maskImage) {
            resizedMask.recycle();
        }

        // 统计输入数据范围，用于调试
        float minVal = Float.MAX_VALUE, maxVal = Float.MIN_VALUE;
        for (float val : inputData) {
            minVal = Math.min(minVal, val);
            maxVal = Math.max(maxVal, val);
        }
        Log.d(TAG, "输入数据范围: [" + minVal + ", " + maxVal + "]");

        return inputData;
    }
    
    /**
     * 转换ONNX输出为float数组
     */
    private float[] convertOutputToFloatArray(Object outputObject) {
        try {
            if (outputObject instanceof float[][][][]) {
                // 4维张量格式 [batch, channel, height, width]
                float[][][][] output4D = (float[][][][]) outputObject;
                
                int batch = output4D.length;
                int channels = output4D[0].length;
                int height = output4D[0][0].length;
                int width = output4D[0][0][0].length;
                
                Log.d(TAG, "输出张量形状: [" + batch + ", " + channels + ", " + height + ", " + width + "]");
                
                // 转换为1维数组
                float[] result = new float[channels * height * width];
                int index = 0;
                for (int c = 0; c < channels; c++) {
                    for (int h = 0; h < height; h++) {
                        for (int w = 0; w < width; w++) {
                            result[index++] = output4D[0][c][h][w];
                        }
                    }
                }
                
                return result;
                
            } else if (outputObject instanceof float[]) {
                return (float[]) outputObject;
            } else {
                throw new RuntimeException("不支持的输出格式: " + outputObject.getClass().getName());
            }
            
        } catch (Exception e) {
            Log.e(TAG, "输出转换失败", e);
            throw new RuntimeException("输出转换失败", e);
        }
    }
    
    /**
     * 后处理输出数据
     * @param outputData 模型输出数据
     * @param width 图像宽度
     * @param height 图像高度
     * @return 处理后的Bitmap
     */
    private Bitmap postprocessOutput(float[] outputData, int width, int height) {
        Bitmap resultBitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
        int[] pixels = new int[width * height];

        int totalPixels = width * height;

        // 检查输出数据长度
        int expectedLength = 3 * totalPixels; // RGB 3通道
        if (outputData.length != expectedLength) {
            Log.w(TAG, "输出数据长度不匹配，期望: " + expectedLength + ", 实际: " + outputData.length);
        }

        // 统计输出值范围，用于调试
        float minVal = Float.MAX_VALUE, maxVal = Float.MIN_VALUE;
        float sumVal = 0.0f;
        for (float val : outputData) {
            minVal = Math.min(minVal, val);
            maxVal = Math.max(maxVal, val);
            sumVal += val;
        }
        float avgVal = sumVal / outputData.length;
        Log.d(TAG, "输出值统计 - 范围: [" + minVal + ", " + maxVal + "], 平均值: " + avgVal);

        // 检查前几个像素的值
        if (outputData.length >= 9) {
            Log.d(TAG, "前3个像素RGB值: " +
                  "R[" + outputData[0] + "," + outputData[1] + "," + outputData[2] + "] " +
                  "G[" + outputData[totalPixels] + "," + outputData[totalPixels+1] + "," + outputData[totalPixels+2] + "] " +
                  "B[" + outputData[2*totalPixels] + "," + outputData[2*totalPixels+1] + "," + outputData[2*totalPixels+2] + "]");
        }

        for (int h = 0; h < height; h++) {
            for (int w = 0; w < width; w++) {
                int pixelIndex = h * width + w;

                // 从NCHW格式提取RGB值
                float r = outputData[pixelIndex]; // R通道
                float g = outputData[totalPixels + pixelIndex]; // G通道
                float b = outputData[2 * totalPixels + pixelIndex]; // B通道

                // LaMa模型输出通常在[-1,1]范围，需要转换到[0,1]
                // 反标准化: value = (value + 1) / 2
                float processedR = (r + 1.0f) / 2.0f;
                float processedG = (g + 1.0f) / 2.0f;
                float processedB = (b + 1.0f) / 2.0f;

                // 确保值在[0,1]范围内
                processedR = Math.max(0.0f, Math.min(1.0f, processedR));
                processedG = Math.max(0.0f, Math.min(1.0f, processedG));
                processedB = Math.max(0.0f, Math.min(1.0f, processedB));

                // 转换为0-255范围
                int red = (int) (processedR * 255);
                int green = (int) (processedG * 255);
                int blue = (int) (processedB * 255);

                pixels[pixelIndex] = Color.argb(255, red, green, blue);
            }
        }

        resultBitmap.setPixels(pixels, 0, width, 0, 0, width, height);

        Log.d(TAG, "后处理完成，生成图像: " + width + "x" + height);
        return resultBitmap;
    }
    
    /**
     * 释放资源
     */
    public void release() {
        try {
            if (session != null) {
                session.close();
                session = null;
            }
            if (env != null) {
                env.close();
                env = null;
            }
            Log.d(TAG, "LaMa ONNX Runtime资源已释放");
        } catch (Exception e) {
            Log.e(TAG, "释放LaMa ONNX Runtime资源失败", e);
        }
    }
    
    /**
     * 检查是否已初始化
     */
    public boolean isInitialized() {
        return session != null && env != null;
    }
}
